import numpy as np
from utils.rank import FrenquencyRank
from utils.frequency_mask import get_frequency_mask

# self frequency
# grid search
def WSA(img, feature_extractor_net, distance_metric, eps=16, gt=None, show=False, save_dir='./res.png'):
    if gt is None:
        gt = feature_extractor_net(img)

    rank = FrenquencyRank(img, gt, 10)

    def get_adv(wave_length, angle):
        angle *= 22.5
        frequency_mask = get_frequency_mask((img.shape[0], img.shape[1]), wave_length, angle) * eps
        perturb = np.dstack([frequency_mask, frequency_mask, frequency_mask])
        adv = img + perturb
        adv[adv > 255] = 255
        adv[adv < 0] = 0
        return adv

    def get_loss(wave_length, angle):
        adv = get_adv(wave_length, angle)
        skeleton = feature_extractor_net(adv)
        loss = distance_metric(
            skeleton,
            gt
        )
        return loss, adv, skeleton

    for wave_length in range(1, 6):
        for angle in range(8):
            loss, adv, skeleton = get_loss(wave_length, angle)
            rank.check(adv, skeleton, loss, { 'wave_length': wave_length, 'angle': angle })

    # rank.draw(save_dir, show)

    optim = rank.best()
    return optim['adv'], optim['skl'], optim['loss'], optim['params']